# -*- coding: utf-8 -*-
#!/usr/bin python

# --------------------------------------
# Non-maximum Suppression(NMS)
# Written by LiPeng
# --------------------------------------

import numpy as np
import sys
reload(sys)
sys.setdefaultencoding('utf-8')

# 单类别的NMS
def single_cls_nms(dets, thresh):
    """
    Pure Python NMS baseline.
    """
    # dets: np.array([[x1, y1, x2, y2, score],[.....], ..., []])
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        # ovr的长度比order长度少一个，所以将所有的下标后移动一位
        order = order[inds + 1]

    return keep

# 多类别NMS
def multi_cls_nms(predicts_dict, thresh):
    for object_name, bbox in predicts_dict.items(): # 对每个类别分别进行NMS
        bbox_array = np.array(bbox, dtype=np.float)
        x1 = bbox_array[:, 0]
        y1 = bbox_array[:, 1]
        x2 = bbox_array[:, 2]
        y2 = bbox_array[:, 3]
        scores = bbox_array[:, 4]
        order = scores.argsort()[::-1]
        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        keep = []
        while order.size() > 0:
            i = order[0]
            keep.append(i) # 保留当前最大的confidence对应的Bbox索引
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])

            w = np.maximum(0.0, xx2 - xx1 + 1)
            h = np.maximum(0.0, yy2 - yy1 + 1)
            inter = w * h
            iou = inter / (areas[i] + areas[order[1:]] - inter)
            inds = np.where(iou <= thresh)[0]
            order =  order[inds + 1]
        bbox = bbox_array[keep]
        predicts_dict[object_name] = bbox.tolist()
    return predicts_dict


if __name__ == '__main__':
    dets = np.array([
                        [204, 102, 358, 250, 0.5],
                        [257, 118, 380, 250, 0.7],
                        [280, 135, 400, 250, 0.6],
                        [255, 118, 360, 235, 0.7]])
    thresh = 0.3

    print single_cls_nms(dets, thresh)
